# 从modelscope库导入必要的模块
# AutoModelForCausalLM: 用于加载因果语言模型
# AutoTokenizer: 用于加载与模型对应的分词器
# snapshot_download: 用于从ModelScope中心下载模型
from modelscope import AutoModelForCausalLM, AutoTokenizer, snapshot_download
import os  # 用于操作系统相关操作

# 设置模型本地存储路径
# 注意：需要替换为你实际的模型路径
model_path = r'D:\work\dev\model-dir\Qwen\Qwen3-4B'

# 备用代码：从ModelScope中心下载模型的方案
# if not os.path.exists(model_dir):
#     os.makedirs(model_dir)  # 如果目录不存在则创建
# model_path = snapshot_download(model_name, cache_dir=model_dir)  # 下载模型到指定路径

print("model_path:", model_path)  # 打印模型路径用于验证

# 加载分词器和模型
# from_pretrained: 从给定路径加载预训练的分词器/模型
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 加载因果语言模型并自动配置：
# - torch_dtype="auto": 自动确定torch数据类型
# - device_map="auto": 自动将模型层分配到可用设备(CPU/GPU)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype="auto",  # 自动选择最佳数据类型
    device_map="auto"     # 自动分配模型到可用设备
)

# 准备模型输入：创建给模型的提示语
prompt = "以表格的形式介绍你自己模型和参数，包括对应需要的gpu内存，功能."
messages = [
    {"role": "user", "content": prompt}  # 用户角色的单条消息
]

# 应用聊天模板格式化消息：
# - tokenize=False: 返回文本而非分词后的输出
# - add_generation_prompt=True: 添加生成提示
# - enable_thinking=True: 启用模型的思考模式(默认开启)
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,              # 返回原始文本
    add_generation_prompt=True,  # 添加生成提示前缀
    enable_thinking=True        # 启用思考模式
)

# 将文本输入转换为模型兼容的格式：
# - return_tensors="pt": 返回PyTorch张量
# - .to(model.device): 将张量移动到模型所在的设备
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# 使用模型生成文本补全：
# - **model_inputs: 解包输入张量
# - max_new_tokens=32768: 生成新token的最大数量
generated_ids = model.generate(
    **model_inputs,        # 传入模型输入
    max_new_tokens=32768  # 控制生成文本长度
)

# 仅提取生成部分(排除输入部分)：
# - [0]: 批次中的第一个序列
# - [len(...):]: 从输入序列末尾开始切片
# - tolist(): 转换为Python列表
output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

# 解析思考内容(<think>标签之间的内容)：
try:
    # 从后向前查找</think>标记(token_id=151668)
    index = len(output_ids) - output_ids[::-1].index(151668)
except ValueError:  # 如果找不到标记
    index = 0  # 将索引设为0(无思考内容)

# 解码输出：
# - output_ids[:index]: 思考内容(</think>之前)
# - output_ids[index:]: 实际生成内容
# - skip_special_tokens=True: 跳过特殊标记如<pad>, </s>
# - strip("\n"): 去除首尾换行符
thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

# 打印思考过程和最终内容
print("thinking content:", thinking_content)  # 模型的推理过程
print("content:", content)  # 最终生成的回答